import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

class EquationSolverNet(nn.Module):
    def __init__(self):
        """
        初始化神经网络的结构。
        
        该方法定义了神经网络的前向传播结构，包括输入层、隐藏层和输出层。
        输入层接收一个10x10的矩阵和一个10维向量，将其展开成110维的向量作为输入。
        输出层有3个神经元，对应3种不同的解题方法。
        """
        super(EquationSolverNet, self).__init__()  # 调用父类的构造方法
        self.fc1 = nn.Linear(10*10 + 10, 128)  # 输入层：10x10矩阵+10维向量展开后大小为110
        self.fc2 = nn.Linear(128, 64)  # 第一个隐藏层，128个神经元
        self.fc3 = nn.Linear(64, 3)  # 输出层：3种方法对应3个类

    def forward(self, x):
        """
        前向传播函数，定义了神经网络的结构和数据的传播方式。

        参数:
        x (Tensor): 输入数据，形状应与fc1层的输入要求一致。

        返回:
        Tensor: 网络的输出，形状由fc3层的输出决定。

        此函数依次通过三个全连接层（fc1, fc2, fc3），并在fc1和fc2之后使用ReLU激活函数。
        """
        # 第一个全连接层，后接ReLU激活
        x = torch.relu(self.fc1(x))
        # 第二个全连接层，后接ReLU激活
        x = torch.relu(self.fc2(x))
        # 第三个全连接层，没有激活函数，直接输出结果
        x = self.fc3(x)
        return x
if __name__ == "__main__":
    # 创建模型
    model = EquationSolverNet()
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 将数据集转换为 DataLoader 以批量处理数据
    dataset = TensorDataset(X, Y)
    train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

    # 开始训练过程
    num_epochs = 10  # 设置训练轮数
    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        # 遍历数据加载器中的批次
        for inputs, labels in train_loader:
            optimizer.zero_grad()  # 清除上一次的梯度信息
            outputs = model(inputs)  # 执行前向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 执行反向传播
            optimizer.step()  # 更新模型参数
            running_loss += loss.item()
        
        # 打印当前轮次的平均损失
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')

    print("训练完毕！")

    # 生成新的方程组
    A_new, b_new = generate_equation_system()

    # 将矩阵 A 和向量 b 展开并转换为张量
    features_new = np.concatenate([A_new.flatten(), b_new])
    features_new = torch.tensor(features_new, dtype=torch.float32).unsqueeze(0)  # 添加 batch 维度

    # 使用模型进行预测
    model.eval()  # 切换到评估模式
    with torch.no_grad():  # 关闭自动梯度计算
        output = model(features_new)
        predicted_method = torch.argmax(output, dim=1).item()

    # 输出预测的最佳解法
    methods = ["Inverse", "LU", "SVD"]
    print(f"模型预测的最优方法是: {methods[predicted_method]}")
